load("@rules_python//python:defs.bzl", "py_library", "py_test")
load("//tools:build_defs.bzl", "py_cpu_gpu_test")

package(
    default_applicable_licenses = ["//:package_license"],
    default_visibility = [
        ":algorithms_packages",
        "//tensorflow_federated/python/learning:learning_users",
    ],
)

package_group(
    name = "algorithms_packages",
    packages = ["//tensorflow_federated/python/learning/algorithms/..."],
)

licenses(["notice"])

py_library(
    name = "algorithms",
    srcs = ["__init__.py"],
    visibility = ["//tensorflow_federated/python/learning:__pkg__"],
    deps = [
        ":fed_avg",
        ":fed_avg_with_optimizer_schedule",
        ":fed_eval",
        ":fed_prox",
        ":fed_recon",
        ":fed_recon_eval",
        ":fed_sgd",
        ":kmeans_clustering",
        ":mime",
        ":personalization_eval",
    ],
)

py_library(
    name = "fed_avg",
    srcs = ["fed_avg.py"],
    deps = [
        "//tensorflow_federated/python/aggregators:factory",
        "//tensorflow_federated/python/aggregators:factory_utils",
        "//tensorflow_federated/python/aggregators:mean",
        "//tensorflow_federated/python/common_libs:py_typecheck",
        "//tensorflow_federated/python/core/environments/tensorflow_frontend:tensorflow_computation",
        "//tensorflow_federated/python/core/impl/types:computation_types",
        "//tensorflow_federated/python/learning:client_weight_lib",
        "//tensorflow_federated/python/learning:loop_builder",
        "//tensorflow_federated/python/learning/metrics:aggregator",
        "//tensorflow_federated/python/learning/metrics:types",
        "//tensorflow_federated/python/learning/models:functional",
        "//tensorflow_federated/python/learning/models:model_weights",
        "//tensorflow_federated/python/learning/models:variable",
        "//tensorflow_federated/python/learning/optimizers:optimizer",
        "//tensorflow_federated/python/learning/optimizers:sgdm",
        "//tensorflow_federated/python/learning/templates:apply_optimizer_finalizer",
        "//tensorflow_federated/python/learning/templates:composers",
        "//tensorflow_federated/python/learning/templates:distributors",
        "//tensorflow_federated/python/learning/templates:learning_process",
        "//tensorflow_federated/python/learning/templates:model_delta_client_work",
    ],
)

py_cpu_gpu_test(
    name = "fed_avg_test",
    timeout = "long",
    srcs = ["fed_avg_test.py"],
    shard_count = 10,
    deps = [
        ":fed_avg",
        "//tensorflow_federated/python/aggregators:factory_utils",
        "//tensorflow_federated/python/core/test:static_assert",
        "//tensorflow_federated/python/learning:loop_builder",
        "//tensorflow_federated/python/learning:model_update_aggregator",
        "//tensorflow_federated/python/learning/metrics:aggregator",
        "//tensorflow_federated/python/learning/models:model_examples",
        "//tensorflow_federated/python/learning/models:test_models",
        "//tensorflow_federated/python/learning/optimizers:sgdm",
    ],
)

py_library(
    name = "fed_avg_with_optimizer_schedule",
    srcs = ["fed_avg_with_optimizer_schedule.py"],
    deps = [
        ":fed_avg",
        "//tensorflow_federated/python/aggregators:factory",
        "//tensorflow_federated/python/aggregators:mean",
        "//tensorflow_federated/python/common_libs:py_typecheck",
        "//tensorflow_federated/python/core/environments/tensorflow_frontend:tensorflow_computation",
        "//tensorflow_federated/python/core/impl/federated_context:federated_computation",
        "//tensorflow_federated/python/core/impl/federated_context:intrinsics",
        "//tensorflow_federated/python/core/impl/types:computation_types",
        "//tensorflow_federated/python/core/impl/types:placements",
        "//tensorflow_federated/python/core/templates:measured_process",
        "//tensorflow_federated/python/learning:client_weight_lib",
        "//tensorflow_federated/python/learning:loop_builder",
        "//tensorflow_federated/python/learning/metrics:aggregator",
        "//tensorflow_federated/python/learning/metrics:types",
        "//tensorflow_federated/python/learning/models:functional",
        "//tensorflow_federated/python/learning/models:model_weights",
        "//tensorflow_federated/python/learning/models:variable",
        "//tensorflow_federated/python/learning/optimizers:optimizer",
        "//tensorflow_federated/python/learning/templates:apply_optimizer_finalizer",
        "//tensorflow_federated/python/learning/templates:client_works",
        "//tensorflow_federated/python/learning/templates:composers",
        "//tensorflow_federated/python/learning/templates:distributors",
        "//tensorflow_federated/python/learning/templates:learning_process",
        "//tensorflow_federated/python/learning/templates:model_delta_client_work",
    ],
)

py_cpu_gpu_test(
    name = "fed_avg_with_optimizer_schedule_test",
    timeout = "long",
    srcs = ["fed_avg_with_optimizer_schedule_test.py"],
    shard_count = 10,
    deps = [
        ":fed_avg_with_optimizer_schedule",
        "//tensorflow_federated/python/core/test:static_assert",
        "//tensorflow_federated/python/learning:loop_builder",
        "//tensorflow_federated/python/learning:model_update_aggregator",
        "//tensorflow_federated/python/learning/metrics:aggregator",
        "//tensorflow_federated/python/learning/models:model_examples",
        "//tensorflow_federated/python/learning/models:test_models",
        "//tensorflow_federated/python/learning/optimizers:sgdm",
    ],
)

py_library(
    name = "fed_prox",
    srcs = ["fed_prox.py"],
    deps = [
        "//tensorflow_federated/python/aggregators:factory",
        "//tensorflow_federated/python/aggregators:factory_utils",
        "//tensorflow_federated/python/aggregators:mean",
        "//tensorflow_federated/python/common_libs:py_typecheck",
        "//tensorflow_federated/python/core/environments/tensorflow_frontend:tensorflow_computation",
        "//tensorflow_federated/python/core/impl/types:computation_types",
        "//tensorflow_federated/python/learning:client_weight_lib",
        "//tensorflow_federated/python/learning:loop_builder",
        "//tensorflow_federated/python/learning/metrics:aggregator",
        "//tensorflow_federated/python/learning/metrics:types",
        "//tensorflow_federated/python/learning/models:functional",
        "//tensorflow_federated/python/learning/models:model_weights",
        "//tensorflow_federated/python/learning/models:variable",
        "//tensorflow_federated/python/learning/optimizers:optimizer",
        "//tensorflow_federated/python/learning/templates:apply_optimizer_finalizer",
        "//tensorflow_federated/python/learning/templates:composers",
        "//tensorflow_federated/python/learning/templates:distributors",
        "//tensorflow_federated/python/learning/templates:learning_process",
        "//tensorflow_federated/python/learning/templates:proximal_client_work",
    ],
)

py_cpu_gpu_test(
    name = "fed_prox_test",
    timeout = "long",
    srcs = ["fed_prox_test.py"],
    shard_count = 10,
    deps = [
        ":fed_prox",
        "//tensorflow_federated/python/aggregators:factory_utils",
        "//tensorflow_federated/python/core/templates:iterative_process",
        "//tensorflow_federated/python/core/test:static_assert",
        "//tensorflow_federated/python/learning:loop_builder",
        "//tensorflow_federated/python/learning:model_update_aggregator",
        "//tensorflow_federated/python/learning/metrics:aggregator",
        "//tensorflow_federated/python/learning/models:model_examples",
        "//tensorflow_federated/python/learning/models:model_weights",
        "//tensorflow_federated/python/learning/models:test_models",
        "//tensorflow_federated/python/learning/optimizers:sgdm",
        "//tensorflow_federated/python/learning/templates:distributors",
    ],
)

py_library(
    name = "fed_recon",
    srcs = ["fed_recon.py"],
    deps = [
        "//tensorflow_federated/python/aggregators:factory",
        "//tensorflow_federated/python/aggregators:factory_utils",
        "//tensorflow_federated/python/aggregators:mean",
        "//tensorflow_federated/python/common_libs:py_typecheck",
        "//tensorflow_federated/python/core/environments/tensorflow_frontend:tensorflow_computation",
        "//tensorflow_federated/python/core/impl/computation:computation_base",
        "//tensorflow_federated/python/core/impl/federated_context:federated_computation",
        "//tensorflow_federated/python/core/impl/federated_context:intrinsics",
        "//tensorflow_federated/python/core/impl/types:computation_types",
        "//tensorflow_federated/python/core/impl/types:placements",
        "//tensorflow_federated/python/core/templates:measured_process",
        "//tensorflow_federated/python/learning:client_weight_lib",
        "//tensorflow_federated/python/learning:tensor_utils",
        "//tensorflow_federated/python/learning/metrics:aggregator",
        "//tensorflow_federated/python/learning/metrics:keras_finalizer",
        "//tensorflow_federated/python/learning/models:reconstruction_model",
        "//tensorflow_federated/python/learning/optimizers:keras_optimizer",
        "//tensorflow_federated/python/learning/optimizers:optimizer",
        "//tensorflow_federated/python/learning/templates:apply_optimizer_finalizer",
        "//tensorflow_federated/python/learning/templates:client_works",
        "//tensorflow_federated/python/learning/templates:composers",
        "//tensorflow_federated/python/learning/templates:distributors",
        "//tensorflow_federated/python/learning/templates:learning_process",
    ],
)

py_test(
    name = "fed_recon_test",
    srcs = ["fed_recon_test.py"],
    # TODO: b/318429579 - Refactor test into unit and integration testing components.
    tags = ["nokokoro"],
    deps = [
        ":fed_recon",
        "//tensorflow_federated/python/aggregators:differential_privacy",
        "//tensorflow_federated/python/aggregators:factory",
        "//tensorflow_federated/python/aggregators:mean",
        "//tensorflow_federated/python/aggregators:robust",
        "//tensorflow_federated/python/aggregators:sum_factory",
        "//tensorflow_federated/python/core/backends/native:execution_contexts",
        "//tensorflow_federated/python/core/environments/tensorflow_frontend:tensorflow_computation",
        "//tensorflow_federated/python/core/impl/federated_context:federated_computation",
        "//tensorflow_federated/python/core/impl/federated_context:intrinsics",
        "//tensorflow_federated/python/core/impl/types:computation_types",
        "//tensorflow_federated/python/core/impl/types:placements",
        "//tensorflow_federated/python/core/templates:aggregation_process",
        "//tensorflow_federated/python/core/templates:iterative_process",
        "//tensorflow_federated/python/core/templates:measured_process",
        "//tensorflow_federated/python/learning:client_weight_lib",
        "//tensorflow_federated/python/learning/metrics:counters",
        "//tensorflow_federated/python/learning/models:model_weights",
        "//tensorflow_federated/python/learning/models:reconstruction_model",
        "//tensorflow_federated/python/learning/optimizers:sgdm",
        "//tensorflow_federated/python/learning/templates:composers",
        "//tensorflow_federated/python/learning/templates:distributors",
    ],
)

py_library(
    name = "fed_recon_eval",
    srcs = ["fed_recon_eval.py"],
    deps = [
        ":fed_recon",
        "//tensorflow_federated/python/aggregators:mean",
        "//tensorflow_federated/python/common_libs:py_typecheck",
        "//tensorflow_federated/python/core/environments/tensorflow_frontend:tensorflow_computation",
        "//tensorflow_federated/python/core/impl/federated_context:federated_computation",
        "//tensorflow_federated/python/core/impl/federated_context:intrinsics",
        "//tensorflow_federated/python/core/impl/types:computation_types",
        "//tensorflow_federated/python/core/impl/types:placements",
        "//tensorflow_federated/python/core/templates:aggregation_process",
        "//tensorflow_federated/python/core/templates:measured_process",
        "//tensorflow_federated/python/learning/metrics:keras_finalizer",
        "//tensorflow_federated/python/learning/metrics:sum_aggregation_factory",
        "//tensorflow_federated/python/learning/models:reconstruction_model",
        "//tensorflow_federated/python/learning/optimizers:keras_optimizer",
        "//tensorflow_federated/python/learning/templates:client_works",
        "//tensorflow_federated/python/learning/templates:composers",
        "//tensorflow_federated/python/learning/templates:distributors",
        "//tensorflow_federated/python/learning/templates:finalizers",
        "//tensorflow_federated/python/learning/templates:learning_process",
    ],
)

py_test(
    name = "fed_recon_eval_test",
    srcs = ["fed_recon_eval_test.py"],
    deps = [
        ":fed_recon_eval",
        "//tensorflow_federated/python/core/backends/native:execution_contexts",
        "//tensorflow_federated/python/core/environments/tensorflow_frontend:tensorflow_computation",
        "//tensorflow_federated/python/core/impl/federated_context:federated_computation",
        "//tensorflow_federated/python/core/impl/federated_context:intrinsics",
        "//tensorflow_federated/python/core/impl/types:computation_types",
        "//tensorflow_federated/python/core/impl/types:placements",
        "//tensorflow_federated/python/core/impl/types:type_test_utils",
        "//tensorflow_federated/python/core/templates:measured_process",
        "//tensorflow_federated/python/learning/metrics:counters",
        "//tensorflow_federated/python/learning/models:reconstruction_model",
        "//tensorflow_federated/python/learning/optimizers:sgdm",
        "//tensorflow_federated/python/learning/templates:composers",
        "//tensorflow_federated/python/learning/templates:distributors",
        "//tensorflow_federated/python/learning/templates:learning_process",
    ],
)

py_library(
    name = "fed_sgd",
    srcs = ["fed_sgd.py"],
    deps = [
        "//tensorflow_federated/python/aggregators:factory",
        "//tensorflow_federated/python/aggregators:mean",
        "//tensorflow_federated/python/common_libs:py_typecheck",
        "//tensorflow_federated/python/core/environments/tensorflow_frontend:tensorflow_computation",
        "//tensorflow_federated/python/core/impl/federated_context:federated_computation",
        "//tensorflow_federated/python/core/impl/federated_context:intrinsics",
        "//tensorflow_federated/python/core/impl/types:computation_types",
        "//tensorflow_federated/python/core/impl/types:placements",
        "//tensorflow_federated/python/core/templates:measured_process",
        "//tensorflow_federated/python/learning:loop_builder",
        "//tensorflow_federated/python/learning:tensor_utils",
        "//tensorflow_federated/python/learning/metrics:aggregator",
        "//tensorflow_federated/python/learning/metrics:types",
        "//tensorflow_federated/python/learning/models:functional",
        "//tensorflow_federated/python/learning/models:model_weights",
        "//tensorflow_federated/python/learning/models:variable",
        "//tensorflow_federated/python/learning/optimizers:optimizer",
        "//tensorflow_federated/python/learning/templates:apply_optimizer_finalizer",
        "//tensorflow_federated/python/learning/templates:client_works",
        "//tensorflow_federated/python/learning/templates:composers",
        "//tensorflow_federated/python/learning/templates:distributors",
        "//tensorflow_federated/python/learning/templates:learning_process",
    ],
)

py_cpu_gpu_test(
    name = "fed_sgd_test",
    timeout = "long",
    srcs = ["fed_sgd_test.py"],
    shard_count = 10,
    deps = [
        ":fed_sgd",
        "//tensorflow_federated/python/core/test:static_assert",
        "//tensorflow_federated/python/learning:loop_builder",
        "//tensorflow_federated/python/learning:model_update_aggregator",
        "//tensorflow_federated/python/learning/metrics:aggregator",
        "//tensorflow_federated/python/learning/models:functional",
        "//tensorflow_federated/python/learning/models:model_examples",
        "//tensorflow_federated/python/learning/models:model_weights",
        "//tensorflow_federated/python/learning/models:test_models",
        "//tensorflow_federated/python/tensorflow_libs:tensorflow_test_utils",
    ],
)

py_library(
    name = "kmeans_clustering",
    srcs = ["kmeans_clustering.py"],
    deps = [
        "//tensorflow_federated/python/aggregators:factory",
        "//tensorflow_federated/python/aggregators:factory_utils",
        "//tensorflow_federated/python/aggregators:sum_factory",
        "//tensorflow_federated/python/core/environments/tensorflow_frontend:tensorflow_computation",
        "//tensorflow_federated/python/core/impl/federated_context:federated_computation",
        "//tensorflow_federated/python/core/impl/federated_context:intrinsics",
        "//tensorflow_federated/python/core/impl/types:computation_types",
        "//tensorflow_federated/python/core/impl/types:placements",
        "//tensorflow_federated/python/core/templates:measured_process",
        "//tensorflow_federated/python/learning/templates:client_works",
        "//tensorflow_federated/python/learning/templates:composers",
        "//tensorflow_federated/python/learning/templates:distributors",
        "//tensorflow_federated/python/learning/templates:finalizers",
        "//tensorflow_federated/python/learning/templates:learning_process",
    ],
)

py_test(
    name = "kmeans_clustering_test",
    srcs = ["kmeans_clustering_test.py"],
    deps = [
        ":kmeans_clustering",
        "//tensorflow_federated/python/core/backends/native:execution_contexts",
        "//tensorflow_federated/python/core/impl/types:computation_types",
        "//tensorflow_federated/python/core/impl/types:placements",
    ],
)

py_library(
    name = "mime",
    srcs = ["mime.py"],
    deps = [
        "//tensorflow_federated/python/aggregators:factory",
        "//tensorflow_federated/python/aggregators:factory_utils",
        "//tensorflow_federated/python/aggregators:mean",
        "//tensorflow_federated/python/common_libs:py_typecheck",
        "//tensorflow_federated/python/core/environments/tensorflow_frontend:tensorflow_computation",
        "//tensorflow_federated/python/core/impl/federated_context:federated_computation",
        "//tensorflow_federated/python/core/impl/federated_context:intrinsics",
        "//tensorflow_federated/python/core/impl/types:computation_types",
        "//tensorflow_federated/python/core/impl/types:placements",
        "//tensorflow_federated/python/core/impl/types:type_conversions",
        "//tensorflow_federated/python/core/templates:measured_process",
        "//tensorflow_federated/python/learning:client_weight_lib",
        "//tensorflow_federated/python/learning:loop_builder",
        "//tensorflow_federated/python/learning:tensor_utils",
        "//tensorflow_federated/python/learning/metrics:aggregator",
        "//tensorflow_federated/python/learning/metrics:types",
        "//tensorflow_federated/python/learning/models:functional",
        "//tensorflow_federated/python/learning/models:model_weights",
        "//tensorflow_federated/python/learning/models:variable",
        "//tensorflow_federated/python/learning/optimizers:optimizer",
        "//tensorflow_federated/python/learning/optimizers:sgdm",
        "//tensorflow_federated/python/learning/templates:apply_optimizer_finalizer",
        "//tensorflow_federated/python/learning/templates:client_works",
        "//tensorflow_federated/python/learning/templates:composers",
        "//tensorflow_federated/python/learning/templates:distributors",
        "//tensorflow_federated/python/learning/templates:learning_process",
    ],
)

py_cpu_gpu_test(
    name = "mime_test",
    timeout = "long",
    srcs = ["mime_test.py"],
    shard_count = 10,
    # TODO: b/318429579 - Refactor test into unit and integration testing components.
    tags = ["nokokoro"],
    deps = [
        ":fed_avg",
        ":mime",
        "//tensorflow_federated/python/aggregators:factory_utils",
        "//tensorflow_federated/python/aggregators:mean",
        "//tensorflow_federated/python/core/backends/native:execution_contexts",
        "//tensorflow_federated/python/core/environments/tensorflow_frontend:tensorflow_computation",
        "//tensorflow_federated/python/core/impl/federated_context:federated_computation",
        "//tensorflow_federated/python/core/impl/federated_context:intrinsics",
        "//tensorflow_federated/python/core/impl/types:computation_types",
        "//tensorflow_federated/python/core/impl/types:placements",
        "//tensorflow_federated/python/core/templates:iterative_process",
        "//tensorflow_federated/python/core/templates:measured_process",
        "//tensorflow_federated/python/core/test:static_assert",
        "//tensorflow_federated/python/learning:client_weight_lib",
        "//tensorflow_federated/python/learning:loop_builder",
        "//tensorflow_federated/python/learning:model_update_aggregator",
        "//tensorflow_federated/python/learning/metrics:aggregator",
        "//tensorflow_federated/python/learning/metrics:counters",
        "//tensorflow_federated/python/learning/models:functional",
        "//tensorflow_federated/python/learning/models:keras_utils",
        "//tensorflow_federated/python/learning/models:model_examples",
        "//tensorflow_federated/python/learning/models:model_weights",
        "//tensorflow_federated/python/learning/models:test_models",
        "//tensorflow_federated/python/learning/optimizers:adagrad",
        "//tensorflow_federated/python/learning/optimizers:adam",
        "//tensorflow_federated/python/learning/optimizers:optimizer",
        "//tensorflow_federated/python/learning/optimizers:rmsprop",
        "//tensorflow_federated/python/learning/optimizers:sgdm",
        "//tensorflow_federated/python/learning/optimizers:yogi",
        "//tensorflow_federated/python/learning/templates:client_works",
        "//tensorflow_federated/python/learning/templates:distributors",
        "//tensorflow_federated/python/tensorflow_libs:tensorflow_test_utils",
    ],
)

py_library(
    name = "fed_eval",
    srcs = ["fed_eval.py"],
    deps = [
        "//tensorflow_federated/python/aggregators:mean",
        "//tensorflow_federated/python/common_libs:py_typecheck",
        "//tensorflow_federated/python/core/environments/tensorflow_frontend:tensorflow_computation",
        "//tensorflow_federated/python/core/impl/computation:computation_base",
        "//tensorflow_federated/python/core/impl/federated_context:federated_computation",
        "//tensorflow_federated/python/core/impl/federated_context:intrinsics",
        "//tensorflow_federated/python/core/impl/types:computation_types",
        "//tensorflow_federated/python/core/impl/types:placements",
        "//tensorflow_federated/python/core/templates:aggregation_process",
        "//tensorflow_federated/python/core/templates:measured_process",
        "//tensorflow_federated/python/learning:loop_builder",
        "//tensorflow_federated/python/learning/metrics:sum_aggregation_factory",
        "//tensorflow_federated/python/learning/models:functional",
        "//tensorflow_federated/python/learning/models:model_weights",
        "//tensorflow_federated/python/learning/models:variable",
        "//tensorflow_federated/python/learning/templates:client_works",
        "//tensorflow_federated/python/learning/templates:composers",
        "//tensorflow_federated/python/learning/templates:distributors",
        "//tensorflow_federated/python/learning/templates:finalizers",
        "//tensorflow_federated/python/learning/templates:learning_process",
    ],
)

py_cpu_gpu_test(
    name = "fed_eval_test",
    srcs = ["fed_eval_test.py"],
    deps = [
        ":fed_eval",
        "//tensorflow_federated/python/core/backends/native:execution_contexts",
        "//tensorflow_federated/python/core/environments/tensorflow_frontend:tensorflow_computation",
        "//tensorflow_federated/python/core/impl/federated_context:federated_computation",
        "//tensorflow_federated/python/core/impl/federated_context:intrinsics",
        "//tensorflow_federated/python/core/impl/types:computation_types",
        "//tensorflow_federated/python/core/impl/types:placements",
        "//tensorflow_federated/python/core/impl/types:type_conversions",
        "//tensorflow_federated/python/core/impl/types:type_test_utils",
        "//tensorflow_federated/python/core/templates:aggregation_process",
        "//tensorflow_federated/python/core/templates:measured_process",
        "//tensorflow_federated/python/learning/metrics:aggregator",
        "//tensorflow_federated/python/learning/metrics:counters",
        "//tensorflow_federated/python/learning/metrics:sum_aggregation_factory",
        "//tensorflow_federated/python/learning/models:functional",
        "//tensorflow_federated/python/learning/models:keras_utils",
        "//tensorflow_federated/python/learning/models:model_examples",
        "//tensorflow_federated/python/learning/models:model_weights",
        "//tensorflow_federated/python/learning/models:variable",
        "//tensorflow_federated/python/learning/templates:composers",
        "//tensorflow_federated/python/learning/templates:distributors",
        "//tensorflow_federated/python/learning/templates:learning_process",
        "//tensorflow_federated/python/tensorflow_libs:tensorflow_test_utils",
    ],
)

py_library(
    name = "personalization_eval",
    srcs = ["personalization_eval.py"],
    deps = [
        "//tensorflow_federated/python/aggregators:sampling",
        "//tensorflow_federated/python/common_libs:py_typecheck",
        "//tensorflow_federated/python/common_libs:structure",
        "//tensorflow_federated/python/core/environments/tensorflow_frontend:tensorflow_computation",
        "//tensorflow_federated/python/core/impl/computation:computation_base",
        "//tensorflow_federated/python/core/impl/federated_context:federated_computation",
        "//tensorflow_federated/python/core/impl/federated_context:intrinsics",
        "//tensorflow_federated/python/core/impl/types:array_shape",
        "//tensorflow_federated/python/core/impl/types:computation_types",
        "//tensorflow_federated/python/core/impl/types:placements",
        "//tensorflow_federated/python/learning/models:model_weights",
        "//tensorflow_federated/python/learning/models:variable",
    ],
)

py_cpu_gpu_test(
    name = "personalization_eval_test",
    srcs = ["personalization_eval_test.py"],
    tags = ["nokokoro"],
    deps = [
        ":personalization_eval",
        "//tensorflow_federated/python/core/backends/native:execution_contexts",
        "//tensorflow_federated/python/core/impl/types:computation_types",
        "//tensorflow_federated/python/learning:loop_builder",
        "//tensorflow_federated/python/learning/models:keras_utils",
        "//tensorflow_federated/python/learning/models:model_examples",
        "//tensorflow_federated/python/learning/models:model_weights",
    ],
)
